# pip install annoy -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn
# pip install gensim==4.2.0 -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host pypi.tuna.tsinghua.edu.cn
# （嘎嘎快）pip install matplotlib -i https://mirrors.aliyun.com/pypi/simple/
from torch import tensor
from sklearn.metrics import f1_score
from datetime import datetime
import time
from collections import Counter
import re
import jieba

import pandas as pd
import time
import numpy as np
from annoy import AnnoyIndex
from tqdm import tqdm
import os
import gensim
from gensim.models import KeyedVectors
from gensim.scripts.glove2word2vec import glove2word2vec
from gensim.models import Word2Vec
from gensim.models.word2vec import LineSentence
import numpy as np
print("gensim.__version__",gensim.__version__) # 4.2.0
import matplotlib.pyplot as plt
plt.rcParams["font.sans-serif"] = ['Simhei']
plt.rcParams["axes.unicode_minus"] = False
import pandas as pd

def search(key_input,top_k):#key_input用户输入的病症描述,top_k返回的信息条数
    data=pd.read_excel("case_information_1800_csv.xlsx")
    #print(data.columns)
    #print(data.head())
    query_list=data['case_name'].values.tolist()#将Excel表中的case_name列装入列表query_list
    answer_list=data['case_detail_symptoms'].values.tolist()#将Excel表中的case_detail_symptoms列装入列表answer_list
    # 使用word2vec之前先进行word2vec的语料库训练 只需要训练一次 就可以
    # with open("word2vec_txt.txt", "a+", encoding='utf-8') as f:
    #     words = []
    #     for i in tqdm(answer_list):
    #         i = "".join(re.findall('[\u4e00-\u9fa5]', str(i)))
    #         i = " ".join(list(jieba.cut(i, cut_all=False)))
    #         f.write(i)
    #         f.write("\n")
    #model = Word2Vec(LineSentence(open('word2vec_txt.txt', 'r', encoding='utf-8')), sg=0, vector_size=64, window=8,min_count=2, workers=4)
    # 模型保存
    # model.save('word2vec.model')
    # 通过模型加载词向量(recommend)

    model_vec = gensim.models.Word2Vec.load('word2vec.model')
    dic = model_vec.wv.index_to_key

    #print(dic)
    #print("---------------------------------------------------------------------------------------------")
    #print("---------------------------------------------------------------------------------------------")
    print(len(dic))
    #print(model_vec.wv['痔疮'])#表示输出‘痔疮’这个词的特征映射结果
    #print(model_vec.wv.most_similar('痔疮', topn=2))#表示输出‘痔疮’这个词最相似（接近）的两个
    #print(query_list[:10])#输出query_list的前十个数据（原Excel表中的的case_name）


    # 为所有的句子建立一个向量库进行索引：
    f=64
    t = AnnoyIndex(f, 'angular')  # Length of item vector that will be indexed
    for index_i,line in enumerate(tqdm(answer_list)):
        # print(line)
        line = "".join(re.findall('[\u4e00-\u9fa5]', str(line)))
        line = jieba.cut(line, cut_all=False)
        temp_vec=np.zeros((1,64))
        for word in line:
            try:
                vec = model_vec.wv[word]  # 从词库里面取出这些词对应的向量
            except KeyError:  # 当在词向量模型中没有这个单词的向量的时候可以设置错误跳过 然后把这个单词的向量设置为全0
                vec = np.zeros((1, 64))
            temp_vec = temp_vec + vec
        t.add_item(index_i, temp_vec[0])

    # 把最新输入的句子变成向量加到 生成新的模型 检索
    # key_input=input("请输入您的症状：")# 脱出，伴有肛门外分泌物、瘙痒、肛门外硬结、皮肤溃疡等相关症状
    key_input= "".join(re.findall('[\u4e00-\u9fa5]', str(key_input)))
    key_input=jieba.cut(key_input, cut_all=False)
    input_vec=np.zeros((1,64))
    idx_input=len(answer_list)+1
    print('idx_input',idx_input)
    for word in key_input:
        # print(word)
        try:
            vec = model_vec.wv[word]  #从词库里面取出这些词对应的向量
            # print(word)
        except KeyError:        #当在词向量模型中没有这个单词的向量的时候可以设置错误跳过 然后把这个单词的向量设置为全0
            vec=np.zeros((1,64))
        input_vec=input_vec+vec
    t.add_item(idx_input, input_vec[0])

    t.build(10)
    t.save('embeedding.ann')
    u = AnnoyIndex(f, 'angular')
    u.load('embeedding.ann')
    temp=u.get_nns_by_item(idx_input,10)
    print(temp)
    res=[]
    for i in temp[1:top_k+1]:
        # print(answer_list[i])
        #res.append(str(query_list[i])+':'+answer_list[i])

        res.append(str(query_list[i]))
    return res


key_input='患者自觉龟头或阴茎部疼痛、灼热及瘙痒感，与衣裤磨擦后加重。由于阴茎头坠胀及衣裤摩擦，有时活动不便'
top_k=5
res=search(key_input,top_k)
for i in res:
    print(i)

